package nmap

import (
	"fmt"
	"github.com/dlclark/regexp2"
	"sort"
	"strconv"
	"strings"
	"time"
)

type TransportProto string

const (
	TCP TransportProto = "TCP"
	UDP TransportProto = "UDP"
)

// ServiceVersion 一方面作为输出，一方面也作为 Match 的原始规则
// ssh 为例, 不准确但能说明意思，这里面除了 ServiceName 确保存在，其他的都是不确定的
// ServiceName: ssh
// Product: openssh
// Version: 1.8.0
// Info: ssl_proto: ssl_1.0, 表示一些额外的信息
// Hostname: 有的服务会返回 hostname 信息
// OS: 有的服务会返回 OS 信息
// DeviceType: 比如 webcam, printer 等信息
// CPE: nmap 返回的 cpe 信息, 可能有多个, 也可能一个没有
type ServiceVersion struct {
	ServiceName string   `json:"service_name"`
	Product     string   `json:"product"`         // p
	Version     string   `json:"version_verbose"` // v
	Info        string   `json:"info"`            // i
	Hostname    string   `json:"hostname"`        // h
	DeviceType  string   `json:"device_type"`     // d
	OS          string   `json:"os"`              // o
	CPE         []string `json:"cpe"`             // cpe
}

type ServiceInfo struct {
	Host   string
	Port   int
	Proto  TransportProto
	SSL    bool
	Banner string
	ServiceVersion
}

type Probe struct {
	Index       int // 仅做区分用
	Name        string
	Proto       TransportProto
	Payload     []byte
	Rarity      int
	Ports       *RangePort // 包含在 ports 规则内，无视稀有度限制，一定使用该探测规则
	SSLPorts    *RangePort // 类似 ports, 只是表明是 ssl
	Fallback    []string   // 相当于 next probes
	WaitTimeout time.Duration
	Matches     []*Match

	// TcpWrapped time.Duration 只在 NULL 中出现，不管了
}

func (p *Probe) String() string {
	return fmt.Sprintf("%s [timeout:%s][rarity:%d]", p.Name, p.WaitTimeout, p.Rarity)
}

type Match struct {
	Soft bool
	Rule *regexp2.Regexp
	ServiceVersion
}

type ProbeParseResult struct {
	Probes       []*Probe
	ExcludePorts *RangePort
}

func (r *ProbeParseResult) ApplyPatch(patch *ProbePatch) error {
	if patch.Name == "*" {
		for _, cp := range r.Probes {
			err := patch.Apply(cp)
			if err != nil {
				return err
			}
		}
		return nil
	}
	for _, cp := range r.Probes {
		if cp.Name == patch.Name {
			return patch.Apply(cp)
		}
	}
	if patch.Mode != PatchModeAdd {
		return fmt.Errorf("can't found the probe, %s", patch.Name)
	}
	p := &Probe{
		Name:  patch.Name,
		Index: len(r.Probes),
	}
	r.Probes = append(r.Probes, p)
	return patch.Apply(p)
}

type RangePort struct {
	singleTCP  []int
	singleUDP  []int
	segmentTCP [][2]int // []{[from, to]}
	segmentUDP [][2]int
}

func (r *RangePort) InSegment(rs [][2]int, port int) bool {
	for _, r := range rs {
		if r[0] <= port && r[1] >= port {
			return true
		}
	}
	return false
}

func (r *RangePort) exist(rs [][2]int, a []int, port int) bool {
	if len(a) != 0 {
		if i := sort.SearchInts(a, port); i < len(a) && a[i] == port {
			return true
		}
	}
	if len(rs) != 0 {
		return r.InSegment(rs, port)
	}
	return false
}

func (r *RangePort) InTCP(port int) bool {
	return r.exist(r.segmentTCP, r.singleTCP, port)
}

func (r *RangePort) AddTCPPorts(ports ...int) {
	r.singleTCP = append(r.singleTCP, ports...)
	r.tidy(r.segmentTCP, r.singleTCP)
}

func (r *RangePort) AddUDPPorts(ports ...int) {
	r.singleUDP = append(r.singleUDP, ports...)
	r.tidy(r.segmentUDP, r.singleUDP)
}

func (r *RangePort) AddTCPSegments(segments ...[2]int) {
	r.segmentTCP = append(r.segmentTCP, segments...)
}

func (r *RangePort) AddUDPSegments(segments ...[2]int) {
	r.segmentUDP = append(r.segmentUDP, segments...)
}

func (r *RangePort) InUDP(port int) bool {
	return r.exist(r.segmentUDP, r.singleUDP, port)
}

func (r *RangePort) tidy(segments [][2]int, single []int) []int {
	sort.Ints(single)
	result := make([]int, 0, len(single))
	var lastVar = -1
	for _, p := range single {
		if p <= 0 {
			continue
		}
		if p == lastVar {
			continue
		}
		lastVar = p
		if r.InSegment(segments, p) {
			continue
		}
		result = append(result, p)
	}
	return result
}

const parsePortErrTpl = "can't parse Port %s"

// -p T:24,29-30,U:443
// T: for TCP, U: for UDP, S: for SCTP, or P: for IP Protocol.
func ParseRangePort(raw string, defaultProto TransportProto) (*RangePort, error) {
	var portType TransportProto
	ret := &RangePort{}
	parts := strings.Split(raw, ",")

	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part[0] == 'S' || part[0] == 'P' {
			continue
		}
		if part[0] == 'T' {
			portType = TCP
			part = part[2:]
		} else if part[0] == 'U' {
			portType = UDP
			part = part[2:]
		} else {
			portType = defaultProto
		}
		if strings.Contains(part, "-") {
			portRange := strings.Split(part, "-")
			if len(portRange) != 2 {
				return nil, fmt.Errorf(parsePortErrTpl, part)
			}
			from, err := strconv.Atoi(strings.TrimSpace(portRange[0]))
			if err != nil {
				return nil, fmt.Errorf(parsePortErrTpl, err)
			}
			to, err := strconv.Atoi(strings.TrimSpace(portRange[1]))
			if err != nil {
				return nil, fmt.Errorf(parsePortErrTpl, err)
			}
			if from > to {
				return nil, fmt.Errorf(parsePortErrTpl, part)
			}
			if portType == TCP {
				if from == to {
					ret.singleTCP = append(ret.singleTCP, from)
				} else {
					ret.segmentTCP = append(ret.segmentTCP, [2]int{from, to})
				}
			} else {
				if from == to {
					ret.singleUDP = append(ret.singleUDP, from)
				} else {
					ret.segmentUDP = append(ret.segmentUDP, [2]int{from, to})
				}
			}
		} else {
			p, err := strconv.Atoi(part)
			if err != nil {
				return nil, fmt.Errorf(parsePortErrTpl, err)
			}
			if portType == TCP {
				ret.singleTCP = append(ret.singleTCP, p)
			} else {
				ret.singleUDP = append(ret.singleUDP, p)
			}
		}
	}

	// 排序，去重，收束
	if len(ret.singleTCP) != 0 {
		ret.singleTCP = ret.tidy(ret.segmentTCP, ret.singleTCP)
	}
	if len(ret.singleUDP) != 0 {
		ret.singleUDP = ret.tidy(ret.segmentUDP, ret.singleUDP)
	}
	return ret, nil
}

// Service nmap 默认端口服务对照
type Service struct {
	Name     string
	Port     int
	Protocol string
}
